import numpy as np
from sklearn.cluster import KMeans


class Welzl:

    def __init__(self, points=None, radius=1):
        self.points = points
        self.radius = radius

    def run(self, points=None):
        """
        基于 Welzl 算法实现的等半径最少圆覆盖算法。即使用固定半径的圆，
        对点集 points 内的所有点进行覆盖，返回最少所需的圆的数目
        :param points: 需要覆盖的全部点集
        :return: 覆盖结果元组，包括圆心集合和点分组集合
        """
        if points is not None:
            self.points = points
        # 1. 初始化参数
        n = len(self.points)
        idx = np.arange(n).reshape((n, 1))
        points = np.concatenate((idx, self.points), axis=1)
        # 2. 迭代处理点集并标记
        circles_pos = []
        points_checked = []
        partition_group = []
        while len(points_checked) < n:
            # 2.1. 缩小有效区域，找出待处理点集中位于区域边界的点
            min_x = np.min(points[:, 1])
            max_x = np.max(points[:, 1])
            min_y = np.min(points[:, 2])
            max_y = np.max(points[:, 2])
            edge = [min_x, min_y, max_x, max_y]
            points_on_edge = [p for p in points if p[1] in edge or p[2] in edge]
            # 2.2. 遍历处理区域边界上的点
            for point_on_edge in points_on_edge:
                # 2.2.1. 当前点已经被处理了（被加入到其它点划分的圆内了）
                point_id = int(point_on_edge[0])
                if point_id in points_checked:
                    continue
                # 2.2.2. 找出所有未处理的点
                mask = np.isin(points[:, 0], np.array(points_checked))
                points_uncheck = points[~mask]
                # 2.2.3. 按照与当前点的距离排序
                points_distance = np.linalg.norm(points_uncheck[:, 1:] - point_on_edge[1:], axis=1)
                points_uncheck_with_dis = np.concatenate((points_uncheck, points_distance[:, np.newaxis]), axis=1)
                sorted_points_uncheck = points_uncheck_with_dis[points_uncheck_with_dis[:, -1].argsort()]
                # 2.2.4. 依次将距离最近的点集加入，调用 Welzl 方法判断最小圆的半径是否超出限制
                #        由于距离最近的点一定是当前点，所以 coverage 方法至少会成功一次
                points_set = []
                points_pos = []
                circle_pos = None
                for point in sorted_points_uncheck:
                    points_set.append(int(point[0]))
                    points_pos.append(np.array([point[1], point[2]]))
                    circle = self.coverage(points_pos)
                    if circle:
                        circle_pos_new, radius = circle
                        if radius > self.radius:
                            points_set = points_set[:-1]
                            break
                        circle_pos = circle_pos_new
                # 2.2.5. 保存以 point_on_edge 为初始点生成的一个固定半径（其实是不超过指定半径）圆覆盖
                circles_pos.append(circle_pos)
                partition_group.append(points_set)
                points_checked.extend(points_set)
                # 2.2.6. 移除已处理的点，以供后续缩小区域边界
                mask = np.isin(points[:, 0], np.array(points_checked))
                points = points[~mask]
        # 3. 返回分区结果
        circles_pos = np.concatenate((circles_pos, np.zeros(len(partition_group))[:, np.newaxis]), axis=1)
        return circles_pos, partition_group

    @classmethod
    def coverage(cls, points):
        """
        针对给定的点集，尝试构建最小圆将所有点都包含在圆上或圆内，
        如果构建成功则返回圆心和半径的二元组，否则返回 None
        :param points: 最小圆覆盖的全部点集
        :return: 构建成功则返回形如 center, radius 的二元组，否则返回 None
        """
        points_num = len(points)
        if points_num == 0:
            return None
        elif points_num == 1:
            return points[0], 0
        elif points_num == 2:
            center = np.mean(points, axis=0)
            radius = np.linalg.norm(points[0] - points[1]) / 2
            return center, radius
        else:
            return cls.welzl(points, [])

    @classmethod
    def welzl(cls, points, boundary):
        """
        给定待考察点集，和边界点集，判断能否完成对所有点的圆覆盖。
        本方法为递归方法，主要流程在 Step.3 首先不断缩小 points
        的规模，并依次尝试扩大所选点集范围，增大覆盖所有 points
        所需的最小圆半径。
        """
        points_num = len(points)
        boundary_num = len(boundary)
        # 1. 边界点集数目达到 3，可唯一确定一个圆
        if boundary_num == 3:
            return cls.generate_circle(boundary[0], boundary[1], boundary[2])
        # 2. 不存在剩余待考察点，且边界点集数目不足 3
        if points_num == 0:
            # 2.1. 边界两点成圆
            if boundary_num == 2:
                center = np.mean(boundary, axis=0)
                radius = np.linalg.norm(boundary[0] - boundary[1]) / 2
                return center, radius
            # 2.2. 边界一点成圆，边界没有点返回 None
            return (boundary[0], 0) if (boundary_num == 1) else None
        # 3. 存在剩余待考察点，递归依次缩小待考察点集规模
        point = points[0]
        circle = cls.welzl(points[1:], boundary)
        if (circle is not None) and cls.point_in_circle(point, circle):
            return circle
        # 4. 递归扩张边集，将外部点加入最小圆内
        new_boundary = boundary + [point]
        circle = cls.welzl(points[1:], new_boundary)
        return circle

    @classmethod
    def generate_circle(cls, a, b, c):
        """
        给定三个点，构造以这三个点为边界的圆，注意三点可唯一确定一个
        圆，且位置和半径均固定。
        """
        d = 2 * (a[0] * (b[1] - c[1]) + b[0] * (c[1] - a[1]) + c[0] * (a[1] - b[1]))
        ux = ((np.linalg.norm(a) ** 2) * (b[1] - c[1]) +
              (np.linalg.norm(b) ** 2) * (c[1] - a[1]) + (np.linalg.norm(c) ** 2) * (a[1] - b[1])) / d
        uy = ((np.linalg.norm(a) ** 2) * (c[0] - b[0]) +
              (np.linalg.norm(b) ** 2) * (a[0] - c[0]) + (np.linalg.norm(c) ** 2) * (b[0] - a[0])) / d
        center = np.array([ux, uy])
        radius = np.linalg.norm(center - a)
        return center, radius

    @classmethod
    def point_in_circle(cls, point, circle):
        center, radius = circle
        return np.linalg.norm(point - center) <= radius


class Cluster:

    def __init__(self, env):
        self.env = env

    def run(self, points):
        g_max = np.ceil((self.env.config.area ** 2) / (2 * (self.env.config.uav_radius ** 2))).astype(int)
        g_left = 1
        g_right = g_max
        print(f"\tUDs={self.env.M}, g_max={g_max}, g={int((g_left + g_right) / 2)}, binary searching the best g")
        # 2. 迭代调用 KMeans，二分法调整分区数目
        kmeans, last_kmeans = None, None
        while g_left <= g_right:
            # 2.1. KMeans
            g = min(int((g_left + g_right) / 2), self.env.M)
            last_kmeans = kmeans
            kmeans = KMeans(n_clusters=g, random_state=0, n_init=10)
            kmeans.fit(points)
            labels = kmeans.labels_
            centers = kmeans.cluster_centers_
            # 2.2. Check partition result
            valid = True
            failed_counter = 0
            for ue_idx, ue_label in enumerate(labels):
                if np.linalg.norm(points[ue_idx] - centers[ue_label]) > self.env.config.uav_radius:
                    valid = False
                    failed_counter += 1
            # 2.3. Binary Search
            if valid:
                print(f"\tg={g}, valid={valid}, g_next=[{g_left}, {g-1}]")
                g_right = g - 1
                last_kmeans = kmeans
            else:
                print(f"\tg={g}, valid={valid}, failed={failed_counter}, g_next=[{g+1}, {g_right}]")
                g_left = g + 1
        # 3. 保存聚类中心点
        g = len(last_kmeans.cluster_centers_)
        parts_ues = []
        for part_idx in range(g):
            ues_idx = [ue_idx for ue_idx, ue_label in enumerate(last_kmeans.labels_) if ue_label == part_idx]
            parts_ues.append(ues_idx)
        return last_kmeans.cluster_centers_, parts_ues
